from functools import wraps
from flask_jwt_extended import jwt_required, get_jwt_identity
from flask import jsonify
from models import Account

def role_required(*required_roles, **kwargs):
    """
    角色要求装饰器
    用法：
    @role_required("ADMIN")  # 单个角色
    @role_required(role=["ADMIN", "MANAGER"])  # 多个角色
    @role_required()  # 仅验证登录

    # 1. 单个角色要求 - 直接传字符串
    @role_required("ADMIN")
    def admin_only():
        return "Only admins can see this"

    # 2. 多个角色要求 - 使用role参数传列表
    @role_required(role=["ADMIN", "MANAGER"])
    def admin_or_manager():
        return "Admins or managers can see this"

    # 3. 只验证用户登录，不检查角色
    @role_required()
    def any_logged_in_user():
        return "Any logged in user can see this"

    # 4. 也支持单个角色作为role参数
    @role_required(role="ADMIN")
    def another_admin_only():
        return "Only admins can see this"    
    """
    def decorator(fn):
        @wraps(fn)
        @jwt_required()
        def wrapper(*args, **kwargs):
            current_user = get_jwt_identity()
            user = Account.query.filter_by(username=current_user).first()
            
            if user is None:
                return jsonify({"msg": "User not found"}), 403
            
            # 获取所有可能的角色要求
            roles_from_args = required_roles  # 元组参数
            roles_from_kwargs = kwargs.get('role', [])  # role关键字参数
            
            # 如果两种方式都没有指定角色，则只验证用户存在
            if not roles_from_args and not roles_from_kwargs:
                return fn(*args, **kwargs)
            
            # 合并所有角色要求到一个列表
            if isinstance(roles_from_kwargs, (list, tuple)):
                required = list(roles_from_args) + list(roles_from_kwargs)
            else:
                required = list(roles_from_args) + [roles_from_kwargs]
            
            # 获取用户的所有启用状态的角色值
            user_role_values = [r.role_value for r in user.roles if r.status == '1']
            
            # 检查是否有任何一个要求的角色
            if not any(role in user_role_values for role in required):
                return jsonify({"msg": "Insufficient role"}), 403
            
            return fn(*args, **kwargs)
            
        return wrapper
    
    # 处理直接调用装饰器的情况
    if len(required_roles) == 1 and callable(required_roles[0]):
        fn = required_roles[0]
        required_roles = ()
        return decorator(fn)
    
    return decorator

# 使用示例：

